import copy

import numpy as np
import torch
from matplotlib import pyplot as plt
from matplotlib.colors import LinearSegmentedColormap
from torch.nn.functional import cosine_similarity


def non_iid_distribution(traindata_cls_counts):
    # 指定点的坐标和大小
    x = []
    y = []
    size = []
    for j in range(len(traindata_cls_counts)):
        for i in range(10):
            try:
                size.append(traindata_cls_counts[j][i]/3)
            except KeyError as e:
                size.append(0)
            x.append(j+1)
            y.append(i+1)

    # 创建图形和轴
    fig, ax = plt.subplots()

    # 在指定的坐标处画点
    # ax.scatter([x1, x2], [y1, y2], s=[size1, size2])
    ax.scatter(x, y, s=size)

    # # 在每个点上添加文字标签
    # for i, txt in enumerate(size):
    #     ax.text(x[i], y[i], str(i + 1), ha='center', va='center')

    # 设置坐标轴的范围，确保点可见
    ax.set_xlim([0, len(traindata_cls_counts) + 1])
    ax.set_ylim([0, len(traindata_cls_counts) + 1])

    # 添加标题、x轴标签和y轴标签
    ax.set_title('α = 0.1')
    ax.set_xlabel('Client ID')
    ax.set_ylabel('Label')

    # 显示图形
    plt.show()


def non_iid_hot(traindata_cls_counts, config):
    # 构建矩阵表示
    matrix = np.zeros((len(traindata_cls_counts), 10))
    for j in range(len(traindata_cls_counts)):
        for i in range(10):
            try:
                matrix[j][i] = traindata_cls_counts[j][i]
            except KeyError as e:
                matrix[j][i] = 0

    # 创建自定义色彩映射
    cmap = LinearSegmentedColormap.from_list("", ["lightblue", "blue"])

    # 创建热力图
    fig, ax = plt.subplots()
    im = ax.imshow(matrix, cmap=cmap)

    # # 显示数值
    # for i in range(len(traindata_cls_counts)):
    #     for j in range(10):
    #         text = ax.text(j, i, int(matrix[i, j]), ha='center', va='center', color='black')

    # 设置x轴刻度
    ax.set_xticks(np.arange(10))
    ax.set_xticklabels(np.arange(1, 11))

    # 设置y轴刻度
    ax.set_yticks(np.arange(len(traindata_cls_counts)))
    ax.set_yticklabels(np.arange(1, len(traindata_cls_counts) + 1))

    # 添加标题、x轴标签和y轴标签
    if config.partition == 'dirichlet':
        ax.set_title(f'α = {config.beta}')
    else:
        ax.set_title('iid')
    ax.set_xlabel('Label')
    ax.set_ylabel('Client ID')

    # 显示颜色条
    plt.colorbar(im, ax=ax)

    # 调整布局以适应数值显示
    plt.tight_layout()

    # 显示图形
    plt.show()


def draw_model(nets):
    # 假设你有四个模型的参数
    models_params = []
    for net in nets:
        models_params.append(copy.deepcopy(net.state_dict()))

    scores = []
    for i in range(len(models_params)):
        line = []
        for j in range(len(models_params)):
            score = cosine_similarity_between_models(models_params[i], models_params[j])
            line.append(score)

        scores.append(line)

    print(scores)
    # 创建热力图
    plt.imshow(scores, cmap='coolwarm', interpolation='nearest')

    # 添加颜色条
    plt.colorbar()

    # 添加坐标轴标签和标题
    plt.xlabel('Client ID')
    plt.ylabel('Client ID')
    plt.title('CIFAR10 iid')
    plt.show()


def cosine_similarity_between_models(model1_params, model2_params):
    # Extract parameter values from state_dict
    model1_params_tensor = torch.cat([v.view(-1) for v in model1_params.values()])
    model2_params_tensor = torch.cat([v.view(-1) for v in model2_params.values()])

    # Calculate cosine similarity matrix
    similarity_matrix = cosine_similarity(model1_params_tensor, model2_params_tensor, dim=0)
    similarity_matrix = similarity_matrix.to('cpu')
    return similarity_matrix.numpy()  # Convert to a single cosine similarity value